# @lhc
# 2025年8月8日

import os
import json
import requests
import zipfile
import logging
import io
from urllib.parse import urlparse

# 配置日志（同时输出到文件和控制台）
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# 文件日志
file_handler = logging.FileHandler('download.log', mode='w')
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

# 控制台日志
console_handler = logging.StreamHandler()
console_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))

logger.addHandler(file_handler)
logger.addHandler(console_handler)

def ensure_dir(directory):
    """确保目录存在"""
    if not os.path.exists(directory):
        os.makedirs(directory)

def process_data_file(data_file):
    """处理数据文件"""
    ensure_dir('image')
    
    with open(data_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            try:
                # 解析每行数据
                parts = [p.strip() for p in line.strip().split('   ') if p.strip()]
                if len(parts) < 5:
                    logger.warning(f"第{line_num}行数据格式不正确，跳过")
                    continue
                
                model_name = parts[0]
                url = parts[4]
                
                logger.info(f"\n{'='*40}\n开始处理: {model_name} | URL: {url}")
                
                # 1. 请求URL获取JSON响应
                try:
                    response = requests.get(url, timeout=10)
                    response.raise_for_status()
                    json_data = response.json()
                    logger.info(f"请求成功 | 响应数据: {json.dumps(json_data, indent=2, ensure_ascii=False)}")
                except Exception as e:
                    logger.error(f"请求失败: {str(e)}")
                    continue
                
                # 2. 获取压缩包URL
                try:
                    zip_url = json_data['data']['url']
                    logger.info(f"获取到压缩包URL: {zip_url}")
                except KeyError:
                    logger.error("响应数据中缺少压缩包URL")
                    continue
                
                # 3. 下载压缩包到内存
                try:
                    logger.info("正在下载压缩包...")
                    zip_response = requests.get(zip_url, stream=True, timeout=30)
                    zip_response.raise_for_status()
                    
                    # 使用内存文件处理
                    with io.BytesIO() as zip_buffer:
                        for chunk in zip_response.iter_content(chunk_size=8192):
                            zip_buffer.write(chunk)
                        zip_buffer.seek(0)
                        
                        # 4. 解压并提取PNG
                        with zipfile.ZipFile(zip_buffer) as zip_ref:
                            png_found = False
                            for file in zip_ref.namelist():
                                if file.lower().endswith('.png'):
                                    # 提取PNG文件
                                    png_data = zip_ref.read(file)
                                    save_path = os.path.join('image', f"{model_name}.png")
                                    with open(save_path, 'wb') as png_file:
                                        png_file.write(png_data)
                                    logger.info(f"✅ 成功提取并保存图片: {save_path}")
                                    png_found = True
                                    break
                            
                            if not png_found:
                                logger.warning("压缩包中未找到PNG文件")
                
                except Exception as e:
                    logger.error(f"下载/解压失败: {str(e)}")
                    continue
                
            except Exception as e:
                logger.error(f"处理第{line_num}行时发生意外错误: {str(e)}")
                continue

if __name__ == '__main__':
    print("="*50)
    print("开始处理数据文件...")
    print("="*50)
    process_data_file('data.txt')
    print("\n处理完成！请查看：")
    print("- image 文件夹（提取的PNG图片）")
    print("- download.log（详细日志）")